/**
 * @file main.cpp
 *
 * Copyright (C) 2024. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include "data_utils.h"
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
extern void eye_custom_do(uint32_t blockDim, void *stream, uint8_t *out,
                          uint32_t batch, uint32_t rows, uint32_t cols, int32_t k);
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__
void eye_custom(GM_ADDR out,
                uint32_t batch,
                uint32_t rows,
                uint32_t cols,
                int32_t k);
#endif

int32_t main(int32_t argc, char *argv[])
{
    (void)argc; (void)argv;
    const uint32_t batch = 8, rows = 1024, cols = 1024;
    const int32_t  k = 0;
    uint32_t blockDim = batch;
    uint64_t perMatrixElems = (uint64_t)rows * cols;
    uint64_t totalElems = (uint64_t)batch * perMatrixElems;
    size_t outputBytes = totalElems * sizeof(float);

#ifdef ASCENDC_CPU_DEBUG
    uint8_t *out = (uint8_t*)AscendC::GmAlloc(outputBytes);
    // 置 0
    for (size_t i = 0; i < outputBytes; ++i) out[i] = 0;
    AscendC::SetKernelMode(KernelMode::AIV_MODE);
    ICPU_RUN_KF(eye_custom, blockDim, out, batch, rows, cols, k);
    WriteFile("./output/output_y.bin", out, outputBytes);
    AscendC::GmFree(out);
#else
    CHECK_ACL(aclInit(nullptr));
    int deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));

    uint8_t *outHost = nullptr;
    uint8_t *outDev  = nullptr;
    CHECK_ACL(aclrtMallocHost((void**)&outHost, outputBytes));
    CHECK_ACL(aclrtMalloc((void**)&outDev, outputBytes, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMemset(outDev, outputBytes, 0, outputBytes));

    eye_custom_do(blockDim, stream, outDev, batch, rows, cols, k);
    CHECK_ACL(aclrtSynchronizeStream(stream));
    CHECK_ACL(aclrtMemcpy(outHost, outputBytes, outDev, outputBytes, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("./output/output_y.bin", outHost, outputBytes);

    CHECK_ACL(aclrtFree(outDev));
    CHECK_ACL(aclrtFreeHost(outHost));
    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    return 0;
}